#include "xtypes.h"
#include "util_iocp.h"


#ifndef _WIN32_WINNT		
#define _WIN32_WINNT 0x0501
#endif
#define WIN32_LEAN_AND_MEAN
#include <windows.h>

#include <WinSock2.h>
#pragma comment( lib, "WS2_32.lib" )
#include <Mswsock.h>
#pragma comment( lib, "Mswsock.lib" )

#include <stddef.h>
#include <stdlib.h>
#include <stdio.h>
#include <conio.h>
#include <string.h>
#include <assert.h>
#include <process.h>



void ResetOverlappedBase(OverlappedBase* base, OverlappedCallback callback)
{
	ZeroMemory(&base->overlapped, sizeof(OVERLAPPED));
	base->callback = callback;
}

void DoOverlappedCallback(OverlappedBase* base, DWORD error, DWORD numberOfBytes, ULONG_PTR key)
{
	assert(base && base->callback);
	
	OverlappedCallback callback = base->callback;

	callback(error, numberOfBytes, key, base);
}



#define COMPLETION_THREAD_COUNT		5

HANDLE	g_completionPort	= NULL;
HANDLE	g_completionThreads[COMPLETION_THREAD_COUNT];

UINT __stdcall CompletionThread(PVOID param)
{
	BOOL			result			= FALSE;
	OverlappedBase* overlapped		= NULL;
	ULONG_PTR		key				= 0;
	DWORD			numberOfBytes	= 0;

	for (;;)
	{
		result = GetQueuedCompletionStatus(
					g_completionPort, 
					&numberOfBytes, 
					&key, 
					(OVERLAPPED**)&overlapped, 
					INFINITE);
		if (result)
		{
			if (numberOfBytes == 0 && key == 0 && !overlapped)
				break;

			OverlappedCallback callback = 
				overlapped->callback;

			callback(
				NO_ERROR, 
				numberOfBytes, 
				key, 
				overlapped);
		}
		else
		{
			if (overlapped)
			{
				OverlappedCallback callback = 
					overlapped->callback;

				if (callback)
				{
					callback(
						GetLastError(), 
						numberOfBytes, 
						key, 
						overlapped);
				}
			}
		}
	}

	return 0;
}

xResult InitializeCompletionPort()
{
	g_completionPort = CreateIoCompletionPort(INVALID_HANDLE_VALUE, NULL, 0, 0);
	if (!g_completionPort)
	{
		return XR_FAIL;
	}

	return XR_OK;
}

xResult ShutdownCompletionPort()
{
	CloseHandle(g_completionPort);
	g_completionPort = NULL;

	return XR_OK;
}

xResult AssociateWithCompletionPort(HANDLE handle, ULONG_PTR key)
{
	if (!g_completionPort)
		return XR_FAIL;

	HANDLE temp = CreateIoCompletionPort(handle, g_completionPort, key, 0);

	if (temp == g_completionPort)
		return XR_OK;

	assert(!"Something went wrong with associating the handle with the completion port!");
	return XR_FAIL;
}

xResult StartCompletionThreads()
{
	xResult ret = XR_OK;

	for (size_t i = 0; i < COMPLETION_THREAD_COUNT; ++i)
	{
		g_completionThreads[i] = (HANDLE)_beginthreadex(NULL, 0, CompletionThread, NULL, 0, NULL);
		if (!g_completionThreads[i])
		{
			ret = XR_FAIL;
			goto Cleanup_Threads;
		}
	}

	return XR_OK;

Cleanup_Threads:
	// Tell the threads that were started, to shut down
	for (size_t i = 0; i < COMPLETION_THREAD_COUNT; ++i)
	{
		if (g_completionThreads[i])
			PostQueuedCompletionStatus(g_completionPort, 0, 0, NULL);
	}

	// Wait for each thread to shutdown
	for (size_t i = 0; i < COMPLETION_THREAD_COUNT; ++i)
	{
		WaitForSingleObject(g_completionThreads[i], INFINITE);
	}

	// Close the handle for each thread
	for (size_t i = 0; i < COMPLETION_THREAD_COUNT; ++i)
	{
		if (g_completionThreads[i])
		{
			CloseHandle(g_completionThreads[i]);
			g_completionThreads[i] = NULL;
		}
	}

	return ret;
}

xResult StopCompletionThreads()
{
	// Tell the threads that were started, to shut down
	for (size_t i = 0; i < COMPLETION_THREAD_COUNT; ++i)
	{
		assert(g_completionThreads[i]);
		PostQueuedCompletionStatus(g_completionPort, 0, 0, NULL);
	}

	// Wait for the threads to shutdown
	WaitForMultipleObjects(
		COMPLETION_THREAD_COUNT, 
		g_completionThreads, 
		TRUE, 
		INFINITE);

	// Close the handle for each thread
	for (size_t i = 0; i < COMPLETION_THREAD_COUNT; ++i)
	{
		CloseHandle(g_completionThreads[i]);
		g_completionThreads[i] = NULL;
	}

	return XR_OK;
}

